import pdfplumber
from typing import Optional
from ..book import Book, Page, Content, TableContent, ContentType
from .exceptions import PageOutOfRangeException
from ..utils import LOG


class PDFParser:
    def __init__(self):
        pass

    def parse_pdf(self, pdf_path: str, start_page: int = 1, end_page: Optional[int] = None) -> Book:
        book = Book(pdf_path)

        with pdfplumber.open(pdf_path) as pdf:
            pdf_pages = len(pdf.pages)
            if end_page is None or end_page > pdf_pages:
                end_page = pdf_pages

            LOG.debug(f"PDF 文件 {pdf_path} 共 {pdf_pages} 页，解析范围: {start_page} - {end_page}")
            if start_page < 0 or start_page > pdf_pages or end_page < start_page:
                raise PageOutOfRangeException(pdf_pages, start_page, end_page)

            page_to_parse = pdf.pages[start_page:end_page]
            for pdf_page in page_to_parse:
                page = Page()

                raw_text = pdf_page.extract_text()
                tables = pdf_page.extract_tables()

                # 将表格数据从文本中移除
                for table_data in tables:
                    for row in table_data:
                        for cell in row:
                            raw_text = raw_text.replace(cell, '', 1)

                if raw_text:
                    raw_text_lines = raw_text.splitlines()
                    cleaned_raw_text_lines = [line.strip() for line in raw_text_lines if line.strip()]
                    cleaned_raw_text = '\n'.join(cleaned_raw_text_lines)

                    text_content = Content(ContentType.TEXT, cleaned_raw_text)
                    page.add_content(text_content)
                    LOG.debug(f"raw_text\n {cleaned_raw_text}")

                if tables:
                    table = TableContent(tables)
                    page.add_content(table)
                    LOG.debug(f"[table]\n {tables}")
                book.add_page(page)

        return book
